import os
from PIL import Image
import random
import argparse
from torchvision import transforms

def split_augment_image(path, file, output, size=1024, num=12):
    image = Image.open(path)
    image_aug = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.RandomCrop(size)
    ])
    _class = path.split("/")[-2]
    filename = file.split(".")
    for i in range(num):
        image_aug(image).save("Photos/{}/{}/{}_{}.{}".format(output, _class, filename[0], str(i).zfill(2), filename[1]))

def makedirs(filename):
    try:
        os.makedirs('Photos' + '/{}'.format(filename))
        os.makedirs('Photos/{}'.format(filename) + '/Benign')
        os.makedirs('Photos/{}'.format(filename) + '/InSitu')
        os.makedirs('Photos/{}'.format(filename) + '/Invasive')
        os.makedirs('Photos/{}'.format(filename) + '/Normal')
    except FileExistsError:
        print("FileExistsError")
        os._exit(0)
    

if __name__=="__main__":
    parser = argparse.ArgumentParser(description='Split image')
    parser.add_argument('--input', '-i', type=str, default='Photos/Normalization/train', help='输入图像路径，文件夹')
    parser.add_argument('--size', '-s', type=int, default=1024, required=True, help='图像分割大小，默认1024')
    parser.add_argument('--num', '-n', type=int, default=12, required=True, help='每张大图增强个数，默认12')
    args = parser.parse_args()
    
    output = "Augmentation_Random_{}_{}".format(args.size, args.num)
    makedirs(output)

    for root, dirs, files in os.walk(args.input):
        for file in files:
            if 'tif' in file:
                split_augment_image("{}/{}".format(root,file), file, output, args.size, args.num)